package com.gitee.wsl.coroutines.atomic

import kotlinx.atomicfu.AtomicLong
import kotlinx.atomicfu.atomic
import kotlinx.atomicfu.locks.SynchronizedObject
import kotlinx.atomicfu.locks.synchronized
import kotlinx.coroutines.CancellableContinuation
import kotlinx.coroutines.currentCoroutineContext
import kotlinx.coroutines.ensureActive
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlin.coroutines.resume

/**
 * Maintains the amount of requests which the client is ready to consume and
 * prevents sending further updates by suspending the sending coroutine
 * if this amount reaches 0.
 *
 * ### Operation
 *
 * Each [useRequest] call decrements the maintained requests amount.
 * Calling coroutine is suspended when this amount reaches 0.
 * The coroutine is resumed when [updateRequests] is called.
 *
 */
class CountLimiter(initial: Int) : SynchronizedObject() {
    private val requests: AtomicLong = atomic(initial.toLong())
    private var awaiter: CancellableContinuation<Unit>? = null

    private fun updateRequests(n: Int) {
        if (n <= 0) return
        synchronized(this) {
            val updatedRequests = requests.value + n.toLong()
            if (updatedRequests < 0) {
                requests.value = Long.MAX_VALUE
            } else {
                requests.value = updatedRequests
            }

            if (awaiter?.isActive == true) {
                awaiter?.resume(Unit)
                awaiter = null
            }
        }
    }

    suspend fun useRequest() {
        if (requests.decrementAndGet() >= 0) {
            currentCoroutineContext().ensureActive()
        } else {
            suspendCancellableCoroutine { continuation ->
                synchronized(this) {
                    if (requests.value >= 0 && continuation.isActive) {
                        continuation.resume(Unit)
                    } else {
                        this.awaiter = continuation
                    }
                }
            }
        }
    }
}